import os
import numpy as np
from astropy.convolution import convolve
from astropy.table import Table
from scipy.ndimage import binary_dilation, median_filter, binary_erosion
from scipy.optimize import minimize
from scipy.stats import trim_mean
from skimage.morphology import remove_small_objects


def xy2vec(xx, yy):
    vec = np.array([xx, yy])
    vec = np.swapaxes(vec, 0, -1)
    return vec


def dist2line(p1, p2, vec):
    v0 = p2 - p1
    return np.cross(v0, vec - p1) / np.linalg.norm(v0)


def weighted_dist(x, *args):
    p1 = np.array([x[0], x[1]])
    if args[2] == 0:
        p2 = np.array([x[0] + 1, x[1] + x[2]])
    else:
        p2 = np.array([x[0] + x[2], x[1] + 1])
    dd = (dist2line(p1, p2, args[0])) ** 2
    flux = args[1]
    return (dd * flux).sum() / flux.sum()


def gauss(x, *p):
    a, mu, sigma = p
    return a * np.exp(-(x - mu) ** 2 / (2. * sigma ** 2))


def get_line_profile(image, return_fit=False):
    ny, nx = image.shape
    yy, xx = np.mgrid[:ny, :nx]

    # line detection
    kernel = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]])
    kernel = kernel.astype(float) / kernel.sum()
    img = image - np.median(image)
    det_img = convolve(img, kernel)
    index = det_img < 0
    sigma = np.sqrt(np.sum(det_img[index] ** 2) / index.sum())
    index = det_img > sigma * 5
    mask = remove_small_objects(index, min_size=3, connectivity=2)

    # extend the mask within the confined region
    xmin, xmax = xx[mask].min(), xx[mask].max()
    ymin, ymax = yy[mask].min(), yy[mask].max()
    index = (xx >= xmin) * (xx <= xmax) * (yy >= ymin) * (yy <= ymax)
    mask = index * binary_dilation(mask, iterations=4)
    xx_good = xx[mask]
    yy_good = yy[mask]
    img_good = img[mask]

    # guess centroid
    x0 = (xx_good * img_good).sum() / img_good.sum()
    y0 = (yy_good * img_good).sum() / img_good.sum()

    # guess the direction of line
    vec_good = xy2vec(xx_good, yy_good)
    vv = vec_good - np.array([x0, y0])
    index = vv[:, 1] < 0
    vv[index, :] = -vv[index, :]  # flip half of the vectors
    va = vv.mean(axis=0)
    if np.abs(va[0]) > 0.1:
        orientation = 0
        rat = va[1] / va[0]
    else:
        orientation = 1
        rat = va[0] / va[1]

    # minimize the sum of weighted distance
    res = minimize(weighted_dist, x0=np.array([x0, y0, rat]),
                   args=(vec_good, img_good, orientation), method='Nelder-Mead')

    # calculate profile
    x1, y1, rat = res.x
    if orientation == 0:
        x2 = x1 + 100
        y2 = y1 + rat * 100
    else:
        x2 = x1 + rat * 100
        y2 = y1 + 100
    p1 = np.array([x1, y1])
    p2 = np.array([x2, y2])
    dd = dist2line(p1, p2, vec_good)

    if return_fit:
        return dd, img_good, p1, p2, mask
    else:
        return dd, img_good


def get_edge_profile(image, smooth=0, return_fit=False):
    ny, nx = image.shape
    yy, xx = np.mgrid[:ny, :nx]
    if smooth > 0:
        data = median_filter(image, smooth * 2 + 1)
    else:
        data = image

    # edge detection
    vmax = np.percentile(data, 95)
    vmin = np.percentile(data, 5)
    thresh = (vmin + vmax) / 2
    index = data > thresh
    edge1 = binary_dilation(index) * (~index)
    edge2 = binary_erosion(index) * (~index)
    edge = np.logical_or(edge1, edge2)
    
    # fit edge
    edge_x, edge_y = xx[edge], yy[edge]
    pfit = np.polyfit(edge_x, edge_y, 1)
    x1 = nx
    y1 = x1 * pfit[0] + pfit[1]
    p1 = np.array([x1, y1])
    x2 = 0
    y2 = pfit[1]
    p2 = np.array([x2, y2])

    # calculate profile
    vec = xy2vec(xx.ravel(), yy.ravel())
    d = dist2line(p1, p2, vec)
    v = data.ravel()
    if v[d > 0].mean() < v[d < 0].mean():
        d = d * (-1)

    if return_fit:
        return d, v, p1, p2, edge
    else:
        return d, v


def grid_curve(dist, value, oversample=5, smooth=0, reject_outlier=True):
    # decide grid position
    dmin, dmax = min(dist), max(dist)
    dmin_outer = int(np.floor(dmin))
    dmax_outer = int(np.floor(dmax)) + 1
    n = (dmax_outer - dmin_outer) * oversample
    xlims = np.arange(n+1) / oversample + dmin_outer
    x = (np.arange(n) + 0.5) / oversample + dmin_outer

    # take average
    y = np.full(n, np.nan)
    for i in range(n):
        index = (dist > xlims[i]) * (dist <= xlims[i+1])
        if np.sum(index) > 0:
            if reject_outlier:
                y[i] = trim_mean(value[index], 0.1)
            else:
                y[i] = np.mean(value[index])

    # fill missing data
    index = np.isfinite(y)
    y = np.interp(x, x[index], y[index], left=np.nan, right=np.nan)
    index = np.isfinite(y)
    x = x[index]
    y = y[index]
    if smooth > 0:
        y = median_filter(y, smooth * 2 + 1)

    return x, y


def get_mtf(value, d=1.0):
    amp = np.absolute(np.fft.fft(value))
    amp = amp / amp[0]
    freq = np.fft.fftfreq(len(amp), d=d)
    index = freq >= 0
    freq = freq[index]
    amp = amp[index]
    return freq, amp


def mtf2sigma(mtf):
    fname = '{}/data/mtf_gaussian.tab'.format(os.path.dirname(__file__))
    tab = Table.read(fname, format='ipac')
    sigma = np.interp(mtf, tab['mtf_ny'], tab['sigma'])
    return sigma
